import torch
import torch.nn.functional as F
from einops import rearrange, repeat
from einops.layers.torch import Rearrange
from torch import einsum, nn

# helpers

def pair(t):
    return t if isinstance(t, tuple) else (t, t)

# classes

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn
    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout = 0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )
    def forward(self, x):
        return self.net(x)

class Attention(nn.Module):
    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
        super().__init__()
        inner_dim = dim_head *  heads
        project_out = not (heads == 1 and dim_head == dim)

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.attend = nn.Softmax(dim = -1)
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity()

    def forward(self, x):
        b, n, _, h = *x.shape, self.heads
        qkv = self.to_qkv(x).chunk(3, dim = -1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
        dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
        attn = self.attend(dots)
        out = einsum('b h i j, b h j d -> b h i d', attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
            ]))
    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x
            x = ff(x) + x
        return x


class Spa_Spe_ViT(nn.Module):
    def __init__(self, *, num_classes=2, dim=81, depth=2, heads=4, mlp_dim=128, pool = 'mean', channels = 3, dim_head = 128, dropout = 0, emb_dropout = 0.):
        super().__init__()
        

        self.to_spectral_embedding = nn.Sequential(
            Rearrange('b c h w -> b c (h w)'),
            # nn.Linear(81, dim),
        )
        
        self.to_spatial_embedding = nn.Sequential(
            Rearrange('b c h w -> b (h w) c'),
            # nn.Linear(60, dim),
        )

        self.pos_spe_embedding = nn.Parameter(torch.randn(1, 60, 81))
        self.pos_spa_embedding = nn.Parameter(torch.randn(1, 81, 60))
        self.pos_embedding = nn.Parameter(torch.randn(1, 60, dim))
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
        self.reality_token = nn.Parameter(torch.randn(1, 1, dim))
        self.dropout = nn.Dropout(emb_dropout)

        self.spectral_transformer = Transformer(81, depth, heads, dim_head, mlp_dim, dropout)
        self.spatial_transformer = Transformer(60, depth, heads, dim_head, mlp_dim, dropout)
        self.pool = pool
        self.to_latent1 = nn.Identity()
        self.to_latent2 = nn.Identity()
        self.mlp_head1 = nn.Sequential(
            # nn.LayerNorm(dim),
            nn.Linear(dim, 2),
            # nn.Sigmoid(),
        )

        self.mlp_head2 = nn.Sequential(
            # nn.LayerNorm(dim),
            nn.Linear(dim, 2),
            # nn.Sigmoid()
        )

    def forward(self, x):
        if x.shape[3] == 60:
            x = x.permute(0,3,1,2)
        # x = self.to_patch_embedding(x) 
        x = x.flatten(2)
        b, n, _ = x.shape
        ################################
        # spectral transformer
        ################################
        x_spectral = x 
        x_spectral = x_spectral + self.pos_spe_embedding
        x_spectral = self.dropout(x_spectral)
        x_spectral = self.spectral_transformer(x_spectral)
        ################################
        # spectral transformer
        ################################


        ################################
        # spatial transformer
        ################################
        x_spatial = x 
        x_spatial = x_spatial.permute(0, 2, 1)
        x_spatial = x_spatial + self.pos_spa_embedding
        x_spatial= self.dropout(x_spatial)
        x_spatial= self.spatial_transformer(x_spatial)
        ################################
        # spatial  transformer
        ################################
        X = x_spectral + x_spatial.permute(0,2,1)
        # cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
        # reality_tokens = repeat(self.reality_token, '() n d -> b n d', b = b)
        # X = torch.cat((cls_tokens, X), dim=1)
        # X = torch.cat((reality_tokens, X), dim=1)

        x1 = X.mean(dim = 1) if self.pool == 'mean' else X[:, 0]
        x2 = X.mean(dim = 1) if self.pool == 'mean' else X[:, 1]
        x_reality = self.to_latent1(x1)
        x_cls = self.to_latent2(x2)
        return self.mlp_head1(x_reality), self.mlp_head2(x_cls)

if __name__ == '__main__':


    img = torch.randn(1, 60, 9, 9)

    preds = v(img) # (1, 1000)
